{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b97c2af7-a217-496b-90df-c3f2b15d0b9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/guodong.li/virtual-venv/peft-venv-py310-cu117/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-07-20 19:48:06,759] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "from peft import get_peft_config, get_peft_model, PromptTuningInit, PromptTuningConfig, TaskType, PeftType\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "import os\n",
    "from transformers import AutoTokenizer\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import default_data_collator, get_linear_schedule_with_warmup\n",
    "from tqdm import tqdm\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "81be8225-e340-4b3e-8126-32e4b035bd48",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\"\n",
    "\n",
    "model_name_or_path = \"/data/nfs/llm/model/bloomz-560m\"\n",
    "tokenizer_name_or_path = \"/data/nfs/llm/model/bloomz-560m\"\n",
    "\n",
    "peft_config = PromptTuningConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    prompt_tuning_init=PromptTuningInit.TEXT,\n",
    "    num_virtual_tokens=8,\n",
    "    prompt_tuning_init_text=\"Classify if the tweet is a complaint or not:\",\n",
    "    tokenizer_name_or_path=model_name_or_path,\n",
    ")\n",
    "\n",
    "dataset_name = \"twitter_complaints\"\n",
    "\n",
    "text_column = \"Tweet text\"\n",
    "label_column = \"text_label\"\n",
    "max_length = 64\n",
    "lr = 3e-2\n",
    "num_epochs = 10\n",
    "batch_size = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5a637672-43a6-4ba4-9905-a070ff3649fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset raft (/home/guodong.li/data/peft/data/raft/twitter_complaints/1.1.0/79c4de1312c1e3730043f7db07179c914f48403101f7124e2fe336f6f54d9f84)\n",
      "100%|██████████| 2/2 [00:00<00:00, 759.01it/s]\n",
      "Loading cached processed dataset at /home/guodong.li/data/peft/data/raft/twitter_complaints/1.1.0/79c4de1312c1e3730043f7db07179c914f48403101f7124e2fe336f6f54d9f84/cache-0e20fff6b1d898ca.arrow\n",
      "Loading cached processed dataset at /home/guodong.li/data/peft/data/raft/twitter_complaints/1.1.0/79c4de1312c1e3730043f7db07179c914f48403101f7124e2fe336f6f54d9f84/cache-8d14a62b8a688c19.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Unlabeled', 'complaint', 'no complaint']\n",
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['Tweet text', 'ID', 'Label', 'text_label'],\n",
      "        num_rows: 50\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['Tweet text', 'ID', 'Label', 'text_label'],\n",
      "        num_rows: 3399\n",
      "    })\n",
      "})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'Tweet text': '@HMRCcustomers No this is my first job',\n",
       " 'ID': 0,\n",
       " 'Label': 2,\n",
       " 'text_label': 'no complaint'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "#dataset = load_dataset(\"ought/raft\", dataset_name)\n",
    "dataset = load_dataset(\"/home/guodong.li/data/peft/raft/raft.py\", dataset_name, cache_dir=\"/home/guodong.li/data/peft/data\")\n",
    "\n",
    "classes = [k.replace(\"_\", \" \") for k in dataset[\"train\"].features[\"Label\"].names]\n",
    "print(classes)\n",
    "\n",
    "dataset = dataset.map(\n",
    "    lambda x: {\"text_label\": [classes[label] for label in x[\"Label\"]]},\n",
    "    batched=True,\n",
    "    num_proc=1,\n",
    ")\n",
    "print(dataset)\n",
    "\n",
    "dataset[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2479040b-0ccf-4100-9429-42737e87ebcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "target_max_length: 3\n",
      "column_names: ['Tweet text', 'ID', 'Label', 'text_label']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch_size: 50\n",
      "0 [227985, 5484, 915, 2566, 169403, 15296, 36272, 525, 3928, 1119, 632, 2670, 3968, 15270, 77658, 915, 210] [1936, 106863, 3]\n",
      "model_inputs input_ids: tensor([     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3, 227985,\n",
      "          5484,    915,   2566, 169403,  15296,  36272,    525,   3928,   1119,\n",
      "           632,   2670,   3968,  15270,  77658,    915,    210,   1936, 106863,\n",
      "             3])\n",
      "model_inputs attention_mask: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
      "labels input_ids: tensor([  -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,   -100,\n",
      "          -100,   -100,   -100,   -100,   -100,   -100,   -100,   1936, 106863,\n",
      "             3])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running tokenizer on dataset:   0%|          | 0/3399 [00:00<?, ? examples/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch_size: 1000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running tokenizer on dataset:  29%|██▉       | 1000/3399 [00:00<00:00, 9029.31 examples/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [227985, 5484, 915, 2566, 74757, 64626, 12384, 44639, 613, 52282, 2670, 79920, 3344, 1002, 368, 17646, 14472, 8348, 664, 718, 4, 19036, 17, 31849, 17, 6312, 76, 44, 62470, 56, 91, 50, 14839, 21, 77658, 915, 210] [3074, 4762, 60943, 3]\n",
      "model_inputs input_ids: tensor([     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3, 227985,   5484,    915,   2566,\n",
      "         74757,  64626,  12384,  44639,    613,  52282,   2670,  79920,   3344,\n",
      "          1002,    368,  17646,  14472,   8348,    664,    718,      4,  19036,\n",
      "            17,  31849,     17,   6312,     76,     44,  62470,     56,     91,\n",
      "            50,  14839,     21,  77658,    915,    210,   3074,   4762,  60943,\n",
      "             3])\n",
      "model_inputs attention_mask: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
      "labels input_ids: tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         3074,  4762, 60943,     3])\n",
      "batch_size: 1000\n",
      "0 [227985, 5484, 915, 232004, 2932, 1188, 221189, 195869, 13, 77658, 915, 210] [3074, 4762, 60943, 3]\n",
      "model_inputs input_ids: tensor([     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3, 227985,   5484,    915, 232004,   2932,   1188,\n",
      "        221189, 195869,     13,  77658,    915,    210,   3074,   4762,  60943,\n",
      "             3])\n",
      "model_inputs attention_mask: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
      "labels input_ids: tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         3074,  4762, 60943,     3])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running tokenizer on dataset:  59%|█████▉    | 2000/3399 [00:00<00:00, 9497.52 examples/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch_size: 1000\n",
      "0 [227985, 5484, 915, 2566, 99198, 53312, 2566, 99198, 53312, 7064, 1074, 1800, 138435, 17, 77658, 915, 210] [3074, 4762, 60943, 3]\n",
      "model_inputs input_ids: tensor([     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3, 227985,   5484,\n",
      "           915,   2566,  99198,  53312,   2566,  99198,  53312,   7064,   1074,\n",
      "          1800, 138435,     17,  77658,    915,    210,   3074,   4762,  60943,\n",
      "             3])\n",
      "model_inputs attention_mask: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
      "labels input_ids: tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         3074,  4762, 60943,     3])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch_size: 399\n",
      "0 [227985, 5484, 915, 2566, 669, 28503, 2633, 34712, 2144, 4867, 2670, 4676, 77658, 915, 210] [3074, 4762, 60943, 3]\n",
      "model_inputs input_ids: tensor([     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "             3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
      "        227985,   5484,    915,   2566,    669,  28503,   2633,  34712,   2144,\n",
      "          4867,   2670,   4676,  77658,    915,    210,   3074,   4762,  60943,\n",
      "             3])\n",
      "model_inputs attention_mask: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,\n",
      "        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])\n",
      "labels input_ids: tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
      "         3074,  4762, 60943,     3])\n",
      "7\n",
      "7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "# data preprocessing\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
    "if tokenizer.pad_token_id is None:\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "\n",
    "target_max_length = max([len(tokenizer(class_label)[\"input_ids\"]) for class_label in classes])\n",
    "print(\"target_max_length:\", target_max_length)\n",
    "\n",
    "\n",
    "# 预处理\n",
    "def preprocess_function(examples):\n",
    "    batch_size = len(examples[text_column])\n",
    "    print(\"batch_size:\", batch_size)\n",
    "    \n",
    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
    "    targets = [str(x) for x in examples[label_column]]\n",
    "    \n",
    "    model_inputs = tokenizer(inputs)\n",
    "    labels = tokenizer(targets)\n",
    "    \n",
    "    for i in range(batch_size):\n",
    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
    "        label_input_ids = labels[\"input_ids\"][i] + [tokenizer.pad_token_id]\n",
    "        if i == 0:\n",
    "            print(i, sample_input_ids, label_input_ids)\n",
    "        model_inputs[\"input_ids\"][i] = sample_input_ids + label_input_ids\n",
    "        labels[\"input_ids\"][i] = [-100] * len(sample_input_ids) + label_input_ids\n",
    "        model_inputs[\"attention_mask\"][i] = [1] * len(model_inputs[\"input_ids\"][i])\n",
    "    #print(model_inputs)\n",
    "    \n",
    "    for i in range(batch_size):\n",
    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
    "        label_input_ids = labels[\"input_ids\"][i]\n",
    "        \n",
    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (max_length - len(sample_input_ids)) + sample_input_ids\n",
    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\"attention_mask\"][i]\n",
    "        labels[\"input_ids\"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids\n",
    "        \n",
    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
    "        labels[\"input_ids\"][i] = torch.tensor(labels[\"input_ids\"][i][:max_length])\n",
    "        if i == 0:\n",
    "            print(\"model_inputs input_ids:\", model_inputs[\"input_ids\"][i])\n",
    "            print(\"model_inputs attention_mask:\", model_inputs[\"attention_mask\"][i])\n",
    "            print(\"labels input_ids:\", labels[\"input_ids\"][i])\n",
    "\n",
    "    \n",
    "        \n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\n",
    "\n",
    "print(\"column_names:\", dataset[\"train\"].column_names)\n",
    "\n",
    "# 将原始的训练和测试数据同时预处理，然后作为训练和评估数据集\n",
    "processed_datasets = dataset.map(\n",
    "    preprocess_function,\n",
    "    batched=True,\n",
    "    num_proc=1,\n",
    "    remove_columns=dataset[\"train\"].column_names,\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Running tokenizer on dataset\",\n",
    ")\n",
    "\n",
    "train_dataset = processed_datasets[\"train\"]\n",
    "eval_dataset = processed_datasets[\"train\"]\n",
    "\n",
    "# 训练与评估使用同一份数据，但是训练数据打乱\n",
    "train_dataloader = DataLoader(train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
    "eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
    "print(len(train_dataloader))\n",
    "print(len(eval_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4d5dd16a-d7c7-480e-98b0-dc7eb23675fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "          227985,   5484,    915,   2566,  74757,  64626,  12384,  44639,    613,\n",
       "           52282,   2670,  79920,   3344,   1002,    368,  17646,  14472,   8348,\n",
       "             664,    718,      4,  19036,     17,  31849,     17,   6312,     76,\n",
       "              44,  62470,     56,     91,     50,  14839,     21,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3, 227985,   5484,    915,    405, 187059,\n",
       "            2256,    664,   2550,  18833,  18607, 162467,      4,   1387,   6199,\n",
       "            3291,  23405,    613,   4657,  17082,    566,   3432,    368,  78851,\n",
       "            1185,  61273,  23181,   1553,  15596,    212, 116057,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3, 227985,   5484,\n",
       "             915,  39762,   2566,  22253,   6201,  75701,     15,    632,    718,\n",
       "            5840,  10006,   6201,  18881,    427,   3804,  19528,    267, 158974,\n",
       "            1320,    368,  10029,    632,  49666,     92,     34,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3, 227985,   5484,    915,   2566, 104565,   8695,   2089,   6140,\n",
       "          109676,  99579,   1369,    512,    368,   4570,     54,    632,    368,\n",
       "            1503, 241485, 132226,     15,    982,    727,   1152,  18100,    861,\n",
       "           32596,  77597, 168154,   1306, 132226,   4346,  87843,     17, 130462,\n",
       "             364,  32923,     89,     53,   8309,     20,     75,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3, 227985,   5484,    915,   2566,\n",
       "           14173,   2960,  29906,    387,  20706,  49337,   1369,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3, 227985,   5484,    915,   2566, 219553,  45736,\n",
       "           36876,   1713,     72,    707, 187205,  13002, 177324,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3, 227985,   5484,    915,   2566, 233938,  28518,  13716,\n",
       "             427,  28146,   1119,  17918,     17, 236706,    368, 214997,   7555,\n",
       "           48659,   5276,  21600,    343,     17,  51416,  22403,    318,   1531,\n",
       "            1306,   1130,  20934,    567, 101161, 184849,  87843,     17,   1594,\n",
       "           15231,   2052,  16642,     20,   7180,     80,     26,  77658,    915,\n",
       "             210],\n",
       "         [     3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "               3,      3,      3,      3,      3,      3,      3,      3,      3,\n",
       "          227985,   5484,    915,   2566,     80,   2068,    479,   2566,     80,\n",
       "            1376,    878, 147587,   3904,    632,    368,   6084,  65673,  78851,\n",
       "           11736,  15527,  19082,  33151,    461,     17,  45575,  17887,    632,\n",
       "            5219,  14216,  68870,   5967,   1841,   4346,  87843,     17,   1594,\n",
       "           14512,     27,     71,   8184,     19,    290,  63748,  77658,    915,\n",
       "             210]]),\n",
       " 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def test_preprocess_function(examples):\n",
    "    batch_size = len(examples[text_column])\n",
    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
    "    model_inputs = tokenizer(inputs)\n",
    "    # print(model_inputs)\n",
    "    for i in range(batch_size):\n",
    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
    "        \n",
    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * ( max_length - len(sample_input_ids)) + sample_input_ids\n",
    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\"attention_mask\"][i]\n",
    "        \n",
    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
    "    return model_inputs\n",
    "\n",
    "# 将原始的测试数据用于测试\n",
    "test_dataset = dataset[\"test\"].map(\n",
    "    test_preprocess_function,\n",
    "    batched=True,\n",
    "    num_proc=1,\n",
    "    remove_columns=dataset[\"train\"].column_names,\n",
    "    load_from_cache_file=False,\n",
    "    desc=\"Running tokenizer on dataset\",\n",
    ")\n",
    "\n",
    "test_dataloader = DataLoader(test_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
    "next(iter(test_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f821e260-fc06-4e0d-8b93-ee9238011384",
   "metadata": {},
   "outputs": [],
   "source": [
    "# creating model\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name_or_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc7f8979-7dcf-440b-93be-f757cb3eafd7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BloomForCausalLM(\n",
       "  (transformer): BloomModel(\n",
       "    (word_embeddings): Embedding(250880, 1024)\n",
       "    (word_embeddings_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    (h): ModuleList(\n",
       "      (0): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (2): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (3): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (4): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (5): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (6): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (7): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (8): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (9): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (10): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (11): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (12): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (13): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (14): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (15): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (16): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (17): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (18): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (19): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (20): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (21): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (22): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (23): BloomBlock(\n",
       "        (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "          (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1024, out_features=250880, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0cfc00f3-2ac3-4ede-b8fd-7356a4873828",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 8,192 || all params: 559,222,784 || trainable%: 0.0014648902430985358\n"
     ]
    }
   ],
   "source": [
    "model = get_peft_model(model, peft_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "622081fb-6478-4c91-8530-6afe5bbf6412",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): BloomForCausalLM(\n",
       "    (transformer): BloomModel(\n",
       "      (word_embeddings): Embedding(250880, 1024)\n",
       "      (word_embeddings_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "      (h): ModuleList(\n",
       "        (0): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (1): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (2): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (3): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (4): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (5): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (6): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (7): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (8): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (9): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (10): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (11): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (12): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (13): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (14): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (15): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (16): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (17): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (18): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (19): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (20): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (21): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (22): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "        (23): BloomBlock(\n",
       "          (input_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (self_attention): BloomAttention(\n",
       "            (query_key_value): Linear(in_features=1024, out_features=3072, bias=True)\n",
       "            (dense): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (post_attention_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): BloomMLP(\n",
       "            (dense_h_to_4h): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "            (gelu_impl): BloomGelu()\n",
       "            (dense_4h_to_h): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln_f): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (lm_head): Linear(in_features=1024, out_features=250880, bias=False)\n",
       "  )\n",
       "  (prompt_encoder): ModuleDict(\n",
       "    (default): PromptEmbedding(\n",
       "      (embedding): Embedding(8, 1024)\n",
       "    )\n",
       "  )\n",
       "  (word_embeddings): Embedding(250880, 1024)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bbf79251-7b7b-4ccd-82d2-341e3f562164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model\n",
    "# optimizer and lr scheduler\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
    "lr_scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=(len(train_dataloader) * num_epochs),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "273c219e-c787-4179-8ff5-5d02cb49e971",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:01<00:00,  4.92it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=0: train_ppl=tensor(2.8311e+13, device='cuda:0') train_epoch_loss=tensor(30.9743, device='cuda:0') eval_ppl=tensor(91053.0859, device='cuda:0') eval_epoch_loss=tensor(11.4192, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.91it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=1: train_ppl=tensor(13540.4053, device='cuda:0') train_epoch_loss=tensor(9.5134, device='cuda:0') eval_ppl=tensor(2038.3656, device='cuda:0') eval_epoch_loss=tensor(7.6199, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.93it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=2: train_ppl=tensor(911.4009, device='cuda:0') train_epoch_loss=tensor(6.8150, device='cuda:0') eval_ppl=tensor(407.7935, device='cuda:0') eval_epoch_loss=tensor(6.0108, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.96it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=3: train_ppl=tensor(362.4507, device='cuda:0') train_epoch_loss=tensor(5.8929, device='cuda:0') eval_ppl=tensor(344.7443, device='cuda:0') eval_epoch_loss=tensor(5.8428, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 11.04it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=4: train_ppl=tensor(280.6489, device='cuda:0') train_epoch_loss=tensor(5.6371, device='cuda:0') eval_ppl=tensor(227.5656, device='cuda:0') eval_epoch_loss=tensor(5.4274, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.96it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=5: train_ppl=tensor(210.2800, device='cuda:0') train_epoch_loss=tensor(5.3484, device='cuda:0') eval_ppl=tensor(201.9129, device='cuda:0') eval_epoch_loss=tensor(5.3078, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.95it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=6: train_ppl=tensor(189.7844, device='cuda:0') train_epoch_loss=tensor(5.2459, device='cuda:0') eval_ppl=tensor(177.7323, device='cuda:0') eval_epoch_loss=tensor(5.1803, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 11.05it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=7: train_ppl=tensor(168.3532, device='cuda:0') train_epoch_loss=tensor(5.1261, device='cuda:0') eval_ppl=tensor(162.9865, device='cuda:0') eval_epoch_loss=tensor(5.0937, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 11.00it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=8: train_ppl=tensor(158.9806, device='cuda:0') train_epoch_loss=tensor(5.0688, device='cuda:0') eval_ppl=tensor(154.2098, device='cuda:0') eval_epoch_loss=tensor(5.0383, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 10.75it/s]\n",
      "100%|██████████| 7/7 [00:00<00:00, 21.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=9: train_ppl=tensor(150.1807, device='cuda:0') train_epoch_loss=tensor(5.0118, device='cuda:0') eval_ppl=tensor(151.4834, device='cuda:0') eval_epoch_loss=tensor(5.0205, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# training and evaluation\n",
    "model = model.to(device)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        #         print(batch)\n",
    "        #         print(batch[\"input_ids\"].shape)\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        total_loss += loss.detach().float()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "    model.eval()\n",
    "    eval_loss = 0\n",
    "    eval_preds = []\n",
    "    for step, batch in enumerate(tqdm(eval_dataloader)):\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        eval_loss += loss.detach().float()\n",
    "        eval_preds.extend(\n",
    "            tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True)\n",
    "        )\n",
    "\n",
    "    eval_epoch_loss = eval_loss / len(eval_dataloader)\n",
    "    eval_ppl = torch.exp(eval_epoch_loss)\n",
    "    train_epoch_loss = total_loss / len(train_dataloader)\n",
    "    train_ppl = torch.exp(train_epoch_loss)\n",
    "    print(f\"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e227b746-fc17-432a-ac7b-d397574fa5a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@TommyHilfiger Dramatic shopping exp. ordered 6 jeans same size (30/32) 2 fits / 2 too large / 2 too slim : same brand &gt; different sizing\n",
      "{'input_ids': tensor([[227985,   5484,    915,   2566, 226154, 126015,   5385,    259, 239364,\n",
      "           3396,  70823,   5853,     17,  57247,   1231, 191040,   5025,   7869,\n",
      "            375,   2324, 149349,     12,    415, 122321,    897,    415,  10136,\n",
      "          10021,    897,    415,  10136,   6497,    381,    915,   5025,  51950,\n",
      "          66869,   5955,    272,  20311,  77658,    915,    210]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
      "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
      "tensor([[227985,   5484,    915,   2566, 226154, 126015,   5385,    259, 239364,\n",
      "           3396,  70823,   5853,     17,  57247,   1231, 191040,   5025,   7869,\n",
      "            375,   2324, 149349,     12,    415, 122321,    897,    415,  10136,\n",
      "          10021,    897,    415,  10136,   6497,    381,    915,   5025,  51950,\n",
      "          66869,   5955,    272,  20311,  77658,    915,    210,   1936, 106863,\n",
      "              2,     31,      4, 115574, 150936,    189,     31,      4]],\n",
      "       device='cuda:0')\n",
      "['Tweet text : @TommyHilfiger Dramatic shopping exp. ordered 6 jeans same size (30/32) 2 fits / 2 too large / 2 too slim : same brand &gt; different sizing Label : no complaint<!DOCTYPE html>\\n<!']\n"
     ]
    }
   ],
   "source": [
    "# 模型评估\n",
    "model.eval()\n",
    "\n",
    "i = 33\n",
    "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
    "print(dataset[\"test\"][i][\"Tweet text\"])\n",
    "print(inputs)\n",
    "\n",
    "with torch.no_grad():\n",
    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "    outputs = model.generate(\n",
    "        input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_new_tokens=10, eos_token_id=3\n",
    "    )\n",
    "    print(outputs)\n",
    "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c062591e-26b2-49d5-8c8a-8a21e772ffe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving model\n",
    "peft_model_id = f\"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\"\n",
    "model.save_pretrained(peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3eeab1ee-4f95-4d9c-9e77-645291109b4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "36K\t/data/nfs/llm/model/bloomz-560m_PROMPT_TUNING_CAUSAL_LM/adapter_model.bin\n",
      "--------------\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "/data/nfs/llm/model/bloomz-560m_PROMPT_TUNING_CAUSAL_LM\n",
      "├── [ 500]  adapter_config.json\n",
      "├── [ 33K]  adapter_model.bin\n",
      "└── [ 129]  README.md\n",
      "\n",
      "0 directories, 3 files\n"
     ]
    }
   ],
   "source": [
    "ckpt = f\"{peft_model_id}/adapter_model.bin\"\n",
    "!du -h $ckpt\n",
    "print(\"--------------\")\n",
    "!tree -h $peft_model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ae29d7c2-9fd5-46ad-b175-2d89090f009d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model path: /data/nfs/llm/model/bloomz-560m\n"
     ]
    }
   ],
   "source": [
    "# 基于高效微调好的模型进行预训练。\n",
    "\n",
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "peft_model_id = f\"{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\"\n",
    "\n",
    "config = PeftConfig.from_pretrained(peft_model_id)\n",
    "print(\"model path:\", config.base_model_name_or_path)\n",
    "# 加载基础模型\n",
    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)\n",
    "# 加载LoRA模型\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7e4fea87-5e32-420b-a1c2-0a9a3a7004e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@greateranglia Ok thanks...\n",
      "{'input_ids': tensor([[227985,   5484,    915,   2566,  14173,   2960,  29906,    387,  20706,\n",
      "          49337,   1369,  77658,    915,    210]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
      "tensor([[227985,   5484,    915,   2566,  14173,   2960,  29906,    387,  20706,\n",
      "          49337,   1369,  77658,    915,    210,   1936, 106863,      2,     31,\n",
      "              4, 115574, 150936,    189,     31,  55683]], device='cuda:0')\n",
      "['Tweet text : @greateranglia Ok thanks... Label : no complaint<!DOCTYPE html>\\n<html>']\n"
     ]
    }
   ],
   "source": [
    "model.to(device)\n",
    "\n",
    "\n",
    "model.eval()\n",
    "i = 4\n",
    "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
    "print(dataset[\"test\"][i][\"Tweet text\"])\n",
    "print(inputs)\n",
    "\n",
    "with torch.no_grad():\n",
    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "    outputs = model.generate(\n",
    "        input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_new_tokens=10, eos_token_id=3\n",
    "    )\n",
    "    print(outputs)\n",
    "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c0c93c-3fa0-42e8-b944-e39b7057775a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
